#!/usr/bin/env python

import optparse
import os
import subprocess
import sys
import tempfile

from charmicat import BadCommandLineError

import ypipi

class Pi0Pi0FitSlurmScript(object):
    script_body = '''#!/bin/bash
#SBATCH -A belle
#SBATCH -t {duration}
#SBATCH -N {nnodes}
#SBATCH -o {stdout}
#SBATCH -e {stderr}

source $HOME/belle/belle_custom_setup.sh
source $HOME/belle/belle_general_setup.sh

ypipi_pi0pi0_fit.py {input_file}
'''

    def __init__(self, input_file, outdir, duration, nnodes, job_name, stdoutdir,
                 stderrdir, stdout, stderr):
        
        output = os.path.join(os.path.abspath(outdir), input_file.rsplit('.', 1)[0] + '.png')
        
        input_file = os.path.abspath(input_file)
        stdoutdir = os.path.abspath(stdoutdir)
        stderrdir = os.path.abspath(stderrdir)
    
        if stdout is None:
            stdout = os.path.join(stderrdir, '{0}_%j.out'.format(job_name))
        
        if stderr is None:
            stderr = os.path.join(stderrdir, '{0}_%j.err'.format(job_name))

        self.script = Pi0Pi0FitSlurmScript.script_body.format(**locals())

    def __str__(self):
        return self.script

        


def main():
    '''
    ypipi_submit_pi0pi0_fits.py [opts] file.h5 [file.h5 [...]]

    Submit via `sbatch' jobs executing `ypipi_pi0pi0_fit.py', one job per input file.
    '''

    op = optparse.OptionParser(usage=main.__doc__)

    op.add_option("-j", "--job-name", dest='job_name', default='pi0pi0_cuts',
                  help='Name jobs NAME in SLURM queue', metavar='NAME')
    op.add_option("-o", "--outdir", dest='outdir', default=os.getcwd(),
                  help='Leave output files in DIR', metavar='DIR')
    op.add_option("-t", "--duration", dest='duration', default='01:00:00',
                  help='Set job duration of N minutes or HH:MM:SS',
                  metavar='N (or) HH:MM:SS')
    op.add_option("--nodes", dest='nnodes', type='int', default=1,
                  help='Use N nodes per job', metavar='N')
    op.add_option("--stdout", dest='stdout', default=None,
                  help='Redirect job stdout to PATH', metavar='PATH')
    op.add_option("--stderr", dest='stderr', default=None,
                  help='Redirect job stderr to PATH', metavar='PATH')
    op.add_option("--stdout-dir", dest='stdoutdir', default=os.getcwd(),
                  help='Place job stdout in DIR', metavar='DIR')
    op.add_option("--stderr-dir", dest='stderrdir', default=os.getcwd(),
                  help='Place job stderr in DIR', metavar='DIR')

    opts, args = op.parse_args()

    if len(args) == 0:
        msg = "Please supply at least one input file."
        raise BadCommandLineError(msg, -1)

    

    for path in args:
        scr = Pi0Pi0FitSlurmScript(os.path.abspath(path), **vars(opts))
        
        with tempfile.NamedTemporaryFile() as fp:
#        with open(path.rsplit('.', 1)[0] + '.sh', 'w') as fp:
            fp.write(str(scr))
            fp.flush()

            sbatch_pipe = subprocess.Popen('sbatch -p shared'.split() + [fp.name], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            
            stdout, stderr = sbatch_pipe.communicate()

            if sbatch_pipe.returncode != 0:
                msg = 'Error submitting job through sbatch:'
                msg += "\n\tSTDOUT: `{0}'".format(stdout)
                msg += "\n\tSTDERR: `{0}'".format(stderr)                
        
                print msg
                print "Continuing..."

                continue

            print stdout.strip()

    
    return 0



if __name__ == "__main__":
    try:
        sys.exit(main())
    except BadCommandLineError, e:
        import os
        sys.stderr.write(str(e) + "\n")
        sys.stderr.write("See `{0} -h' for more.\n".format(os.path.basename(sys.argv[0])))
        sys.stderr.flush()
        sys.exit(e.retcode)
